# application/extensions/init_databases.py (最终版本)
import sys
from flask import Flask
from flask_sqlalchemy import SQLAlchemy
import logging

from application.extensions.logging import system_logger

# 创建数据库实例
db = SQLAlchemy()


def init_databases(app: Flask):
    """初始化数据库"""

    # 初始化数据库
    db.init_app(app)
    
    # 设置SQL日志着色
    setup_sql_logging(app)

    system_logger.info(f"数据库连接成功: {app.config.get('SQLALCHEMY_DATABASE_URI')}")
    return db


def setup_sql_logging(app: Flask):
    """设置SQL日志着色显示"""
    # 只在启用了SQLALCHEMY_ECHO时配置
    if not app.config.get('SQLALCHEMY_ECHO'):
        return
        
    class RedSQLFormatter(logging.Formatter):
        RED = '\033[31m'
        RESET = '\033[0m'
        
        def format(self, record):
            # 为SQL语句添加红色高亮
            if any(keyword in record.getMessage() for keyword in 
                   ['SELECT', 'INSERT', 'UPDATE', 'DELETE', 'BEGIN', 'COMMIT']):
                record.msg = f"{self.RED}{record.msg}{self.RESET}"
            return super().format(record)
    
    # 获取SQLAlchemy的engine logger
    sql_logger = logging.getLogger('sqlalchemy.engine')
    sql_logger.setLevel(logging.INFO)
    
    # 清除现有处理器以避免重复
    sql_logger.handlers.clear()
    
    # 创建控制台处理器专门用于SQL
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(RedSQLFormatter('%(message)s'))
    
    sql_logger.addHandler(console_handler)
    sql_logger.propagate = False